"""
@FileName：main.py
@Description：
@Author：Lemon
@Time：2025/5/8 15:22
"""
import pandas as pd

from writer import MyWriter
import os
from dataset import MyDataset
from torch.utils.data import DataLoader
from train import train
import torch

os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

batch_size = 32
lr = 0.001
epochs = 200
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_dir = './model'
log_dir = './log'
os.makedirs(log_dir, exist_ok=True)
os.makedirs(model_dir, exist_ok=True)
writer = MyWriter(log_dir)


def load_label_map(csv_path):
    df = pd.read_csv(csv_path)
    return {row['file']: row['label'] for _, row in df.iterrows()}


label_map = load_label_map("./processed/label.csv")

dataset_train = MyDataset('./processed/train/', label_map)
train_loader = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True, pin_memory=True,
                          drop_last=True, num_workers=4)

dataset_test = MyDataset('./processed/validate/', label_map)
test_loader = DataLoader(dataset=dataset_test, batch_size=batch_size, shuffle=True, pin_memory=True, drop_last=True,
                         num_workers=4)

if __name__ == '__main__':
    train(train_loader, test_loader, writer, epochs, lr, device, model_dir)

